import matplotlib
import networkx as nx
import matplotlib.pyplot as plt
import json
from networkx.readwrite import json_graph



# 修正中文显示问题
matplotlib.rcParams['font.sans-serif'] = ['Microsoft YaHei']  # 中文字体
matplotlib.rcParams['axes.unicode_minus'] = False  # 处理负号问题
def generate_hierarchical_topology(num_nodes, edges_per_new_node):
    """
    使用 Barabási–Albert 模型生成具有层次结构的随机网络拓扑
    :param num_nodes: 总节点数
    :param initial_nodes: 初始完全连接节点数
    :param edges_per_new_node: 每个新节点的边连接数量
    :return: 生成的图对象
    """
    # 生成 Barabási–Albert 模型图
    G = nx.barabasi_albert_graph(num_nodes, edges_per_new_node)

    # 添加节点属性
    for node in G.nodes():
        G.nodes[node]['type'] = 'center' if len(list(G.neighbors(node))) > edges_per_new_node else 'leaf'

    # 添加边属性
    for u, v in G.edges():
        G.edges[u, v]['weight'] = 1  # 可以根据需要调整权重

    return G


def visualize_topology(G, title="Barabási–Albert层次结构随机网络拓扑图"):
    """
    可视化网络拓扑
    :param G: 图对象
    :param title: 图的标题
    """
    pos = nx.spring_layout(G)  # 使用 spring 布局
    node_colors = ['red' if G.nodes[node]['type'] == 'center' else 'lightblue' for node in G.nodes()]
    nx.draw_networkx_nodes(G, pos, node_color=node_colors, node_size=500)
    nx.draw_networkx_edges(G, pos, edge_color='gray')
    nx.draw_networkx_labels(G, pos, font_size=10, font_color='black')
    plt.title(title)
    plt.show(block=True)


def export_topology_data(G, file_name="hierarchical_graph_data.json"):
    """
    导出拓扑数据到 JSON 文件
    :param G: 图对象
    :param file_name: 输出文件名
    """
    data = json_graph.node_link_data(G)
    with open(file_name, 'w') as f:
        json.dump(data, f, indent=4)
    print(f"拓扑数据已导出到 {file_name}")


def main():
    # 参数设置
    num_nodes = 64  # 总节点数
    # initial_nodes = 5  # 初始完全连接节点数
    edges_per_new_node = 1  # 每个新节点的边连接数量

    # 生成层次结构拓扑
    G = generate_hierarchical_topology(num_nodes, edges_per_new_node)

    # 可视化拓扑
    visualize_topology(G)

    # 导出数据
    export_topology_data(G, "hierarchical_graph_data.json")


if __name__ == "__main__":
    main()
